package com.hwtx.soa.registry.impl;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;

import com.hwtx.soa.registry.ChildListener;
import com.hwtx.soa.registry.StateListener;
import com.hwtx.soa.registry.ZkClient;
import com.hwtx.soa.registry.path.Path;
import com.hzgj.bcl.soa.config.ReferenceDef;
import com.hzgj.bcl.soa.config.ServerDef;
import com.hzgj.bcl.soa.config.ServiceDef;
import com.hzgj.bcl.soa.registry.AbstractRegistry;
import com.hzgj.bcl.soa.registry.NotifyListener;
import com.hzgj.bcl.soa.registry.RegistryContext;
import com.hzgj.bcl.util.MapUtil;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Collection;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

import static com.hzgj.bcl.util.lang.ThreadPoolKit.getScheduledExecutorService;

public class ZkRegistry extends AbstractRegistry {

    private Logger logger = LoggerFactory.getLogger(getClass());

    ZkClient client;

    private Path path;

    private final Set<ServiceDef.Attr> failedRegistered = Sets.newConcurrentHashSet();
    private final ConcurrentMap<ReferenceDef, Set<NotifyListener>> failedSubscribed = Maps.newConcurrentMap();
    private final ConcurrentMap<ReferenceDef, Set<NotifyListener>> failedUnsubscribed = Maps.newConcurrentMap();
    private final Set<ReferenceDef> failedRegistryClients = Sets.newConcurrentHashSet();
    private final ConcurrentMap<NotifyListener, ChildListener> zkListeners = Maps.newConcurrentMap();

    public ZkRegistry(ZkClient client, Path path) {
        this.client = client;
        this.path = path;
        this.client.addStateListener(state -> {
            if (state == StateListener.RECONNECTED) {
                try {
                    recover();
                } catch (Exception e) {
                    logger.error(e.getMessage(), e);
                }
            }
        });

        ScheduledExecutorService scheduledExecutorService = getScheduledExecutorService("zk-retry-service-");
        scheduledExecutorService.scheduleAtFixedRate(() -> {
            try {
                retry();
            } catch (Throwable t) { // Defensive fault tolerance
                logger.error("Unexpected error occur at failed retry, cause: " + t.getMessage(), t);
            }
        }, 0, 2, TimeUnit.SECONDS);
    }

    protected void retry() {
        if (!failedRegistered.isEmpty()) {
            Set<ServiceDef.Attr> failed = new HashSet<>(failedRegistered);
            if (failed.size() > 0) {
                if (logger.isInfoEnabled()) {
                    logger.info("Retry register " + failed);
                }
                try {
                    failed.forEach(attr -> {
                        try {
                            registryService(attr);
                        } catch (Throwable t) { // Ignore all the exceptions and wait for the next retry
                            logger.warn("Failed to retry register " + attr + ", waiting for again, cause: " + t.getMessage(), t);
                        }
                    });
                } catch (Throwable t) { // Ignore all the exceptions and wait for the next retry
                    logger.warn("Failed to retry register " + failed + ", waiting for again, cause: " + t.getMessage(), t);
                }
            }
        }

        if (!failedRegistryClients.isEmpty()) {
            Set<ReferenceDef> failed = new HashSet<>(failedRegistryClients);
            if (failed.size() > 0) {
                if (logger.isInfoEnabled()) {
                    logger.info("Retry register client" + failed);
                }
                try {
                    failed.forEach(ref -> {
                        try {
                            registryClient(ref);
                        } catch (Throwable t) { // Ignore all the exceptions and wait for the next retry
                            logger.warn("Failed to retry register client" + ref + ", waiting for again, cause: " + t.getMessage(), t);
                        }
                    });
                } catch (Throwable t) { // Ignore all the exceptions and wait for the next retry
                    logger.warn("Failed to retry register client" + failed + ", waiting for again, cause: " + t.getMessage(), t);
                }
            }
        }

        if (!failedSubscribed.isEmpty()) {
            Map<ReferenceDef, Set<NotifyListener>> failed = Maps.newHashMap(failedSubscribed);
            if (failed.size() > 0) {
                if (logger.isInfoEnabled()) {
                    logger.info("Retry subscribe " + failed);
                }
                try {
                    failed.forEach((ref, listeners) -> {
                        try {
                            listeners.forEach(listener -> {
                                watchService(listener, ref);
                            });
                        } catch (Throwable t) { // Ignore all the exceptions and wait for the next retry
                            logger.warn("Failed to retry subscribe " + ref + ", waiting for again, cause: " + t.getMessage(), t);
                        }
                    });
                } catch (Throwable t) { // Ignore all the exceptions and wait for the next retry
                    logger.warn("Failed to retry subscribe" + failed + ", waiting for again, cause: " + t.getMessage(), t);
                }
            }
        }
        if (!failedUnsubscribed.isEmpty()) {
            Map<ReferenceDef, Set<NotifyListener>> failed = Maps.newHashMap(failedUnsubscribed);
            if (failed.size() > 0) {
                if (logger.isInfoEnabled()) {
                    logger.info("Retry unsubscribe " + failed);
                }
                try {
                    failed.forEach((referenceDef, listeners) -> {
                        try {
                            unWatchService(listeners, referenceDef);
                        } catch (Throwable t) { // Ignore all the exceptions and wait for the next retry
                            logger.warn("Failed to retry unsubscribe " + failed + ", waiting for again, cause: " + t.getMessage(), t);
                        }

                    });
                } catch (Throwable t) { // Ignore all the exceptions and wait for the next retry
                    logger.warn("Failed to retry unsubscribe " + failed + ", waiting for again, cause: " + t.getMessage(), t);
                }
            }
        }
    }

    protected void recover() {
        // register
        List<ServiceDef.Attr> recoverRegistered = Lists.newArrayList(registered);
        if (!recoverRegistered.isEmpty()) {
            if (logger.isInfoEnabled()) {
                logger.info("Recover register services " + recoverRegistered);
            }
            recoverRegistered.forEach(attr -> {
                failedRegistered.add(attr);
            });
            registered.clear();
        }
        // subscribe
        Map<ReferenceDef, Set<NotifyListener>> recoverSubscribed = Maps.newHashMap(getSubscribed());
        if (!recoverSubscribed.isEmpty()) {
            if (logger.isInfoEnabled()) {
                logger.info("Recover subscribe " + recoverSubscribed.values());
            }
            recoverSubscribed.forEach((ref, listeners) -> {
                failedRegistryClients.add(ref);
                failedSubscribed.put(ref, listeners);
            });
        }
    }

    @Override
    public void registryService(RegistryContext.ServerRegisterModel serverRegisterModel, boolean refreshTtl) {

        List<ServiceDef.Attr> services = serverRegisterModel.getServices().values().stream().filter(attr -> attr.isRegister()).collect(Collectors.toList());
        if (services.size() == 0) {
            if (logger.isDebugEnabled()) {
                logger.debug("server {} has no service to register.", serverRegisterModel.getServerId());
            }
            return;
        }
        registryServices(services);
    }

    private void registryServices(List<ServiceDef.Attr> services) {
        services.forEach(this::registryService);
    }

    private void registryService(ServiceDef.Attr attr) {

        failedRegistered.remove(attr);
        String _path = path.getProviderPath(attr);
        try {
            client.create(_path, true);
            registered.add(attr);
        } catch (Exception e) {
            Throwable t = e;
            logger.error("Failed to register " + _path + ", waiting for retry, cause: " + t.getMessage(), t);
            failedRegistered.add(attr);
        }
    }

    @Override
    public ServiceDef lookupService(String name) {
        ServiceDef serviceDef = service_cache.get(name);
        if (serviceDef == null) {
            getServers(name, client.getChildren(path.getProviderPrefix(name)));
            serviceDef = service_cache.get(name);
        }
        return serviceDef;
    }


    @Override
    public void notifyClient(String service) {
        ReferenceDef referenceDef = new ReferenceDef();
        referenceDef.setInterface(service);
        Set<NotifyListener> listeners = getSubscribed().get(referenceDef);
        if (listeners != null) {
            listeners.forEach(listener -> {
                watchService(listener, referenceDef);
            });
        }
        registryClient(referenceDef);
    }

    Map<String, ServiceDef> service_cache = Maps.newHashMap();

    @Override
    public boolean needLoopPullService() {
        return false;
    }

    @Override
    public void updateClient(Collection<ReferenceDef> referenceDefs) {
        referenceDefs.forEach(referenceDef -> {
            String child = path.getConsumerChildByServiceAndIp(client, referenceDef.getService(), referenceDef.getHost());
            if (child != null) {
                client.create(path.getConsumerPath(referenceDef), true);
                client.delete(path.getConsumerPath(referenceDef.getService(), child));
            }
        });
    }

    @Override
    protected void unWatchService(Set<NotifyListener> listeners, ReferenceDef referenceDef) {

        Set<NotifyListener> _listeners = failedSubscribed.get(referenceDef);
        if (_listeners != null) {
            _listeners.removeAll(listeners);
            if (_listeners.size() == 0) {
                failedSubscribed.remove(referenceDef);
            }
        }

        _listeners = failedUnsubscribed.get(referenceDef);
        if (_listeners != null) {
            _listeners.removeAll(listeners);
            if (_listeners.size() == 0) {
                failedUnsubscribed.remove(referenceDef);
            }
        }

        try {
            if (listeners != null) {
                listeners.forEach(notifyListener -> {
                    ChildListener zkListener = zkListeners.get(notifyListener);
                    if (zkListener != null) {
                        client.removeChildListener(path.getProviderPrefix(referenceDef.getService()), zkListener);
                        zkListeners.remove(notifyListener);
                    }
                });
            }
        } catch (Exception e) {
            Throwable t = e;
            logger.error("Failed to unsubscribe " + referenceDef + ", waiting for retry, cause: " + t.getMessage(), t);
            failedUnsubscribed.put(referenceDef, listeners);
        }
    }

    @Override
    protected void watchService(NotifyListener notifyListener, ReferenceDef referenceDef) {


        Set<NotifyListener> listeners = failedSubscribed.get(referenceDef);
        if (listeners != null) {
            listeners.remove(notifyListener);
            if (listeners.size() == 0) {
                failedSubscribed.remove(referenceDef);
            }
        }
        listeners = failedUnsubscribed.get(referenceDef);
        if (listeners != null) {
            listeners.remove(notifyListener);
            if (listeners.size() == 0) {
                failedUnsubscribed.remove(referenceDef);
            }
        }

        String url = path.getProviderPrefix(referenceDef.getService());
        try {

            ChildListener zkListener = zkListeners.get(notifyListener);
            if (zkListener == null) {
                zkListeners.putIfAbsent(notifyListener, (parentPath, currentChildren) -> {

                });
                zkListener = zkListeners.get(notifyListener);
            }

            List<String> children = client.addChildListener(url, zkListener);
            if (children != null) {
                notify(referenceDef.getService(), notifyListener, children);
            }

        } catch (Exception e) {
            Throwable t = e;
            logger.error("Failed to subscribe " + url + ", waiting for retry, cause: " + t.getMessage(), t);
            MapUtil.putMapValue(failedSubscribed, referenceDef, notifyListener);
        }
    }

    private void notify(String service, NotifyListener listener, List<String> children) {

        List<ServerDef> servers = getServers(service, children);

        if (logger.isDebugEnabled()) {
            logger.debug("notify service " + service);
        }

        if (listener != null) {
            if (servers.size() > 0) {
                listener.notifyWithSingleService(servers);
            } else {
                logger.warn("update server size is 0,client subscribe service is " + service);
            }
        }
    }

    private List<ServerDef> getServers(String service, List<String> children) {
        ServiceDef serviceDef = new ServiceDef();
        serviceDef.setName(service);
        List<ServerDef> servers = Lists.newArrayList();

        children.forEach(child -> {
            ServiceDef.Attr newAttr = path.buildService(child);
            newAttr.setInterface(service);
            serviceDef.getAttrs().add(newAttr);
            servers.add(newAttr.getServerDef());
        });

        if (serviceDef.getAttrs().size() == 0) {
            service_cache.remove(service);
        } else {
            service_cache.put(service, serviceDef);
        }

        return servers;
    }

    @Override
    protected void registryClient(ReferenceDef ref) {
        String _path = path.getConsumerPath(ref);
        try {
            failedRegistryClients.remove(ref);
            client.create(_path, true);
        } catch (Exception e) {
            Throwable t = e;
            logger.error("Failed to register client " + _path + ", waiting for retry, cause: " + t.getMessage(), t);
            failedRegistryClients.add(ref);
        }
    }
}
